import os
import shutil

import numpy as np

from option import parse_args

LIST_OF_TESTING_VIDEOS = 'List_of_testing_videos.txt'


def first_five_frames_exist(real_video_crop, fake_video_crop):
    for i in range(5):
        image_name = '{}_0.png'.format(i * 10)
        file_exist = os.path.exists(os.path.join(real_video_crop, image_name)) and \
                     os.path.exists(os.path.join(fake_video_crop, image_name))
        if not file_exist:
            return False
    return True


def copy_five_frame(sampled_videos, args):
    with open(os.path.join(args.dst_dir, 'real_fake_pairs.dat'), 'w') as f:
        for sampled_video in sampled_videos:
            sampled_video_split = sampled_video.split('_')
            assert len(sampled_video_split) == 3
            real_video_crop = os.path.join(args.src_dir, 'crops', '{}_{}'.format(sampled_video_split[0],
                                                                                 sampled_video_split[2][:-4]))
            fake_video_crop = os.path.join(args.src_dir, 'crops', sampled_video[:-4])
            # real, fake = sampled_video_split[0], sampled_video[:-4]
            if not first_five_frames_exist(real_video_crop, fake_video_crop):
                continue
            for i in range(5):
                image_name = '{}_0.png'.format(i * 10)
                if not os.path.exists(os.path.join(real_video_crop, image_name)):
                    continue

                dst_real_image = f'{sampled_video_split[0]}_{sampled_video_split[2][:-4]}' + '__' + image_name
                if os.path.exists(os.path.join(args.src_dir, dst_real_image)):
                    continue
                dst_fake_image = f'{sampled_video_split[0]}_{sampled_video_split[2][:-4]}_{sampled_video_split[1]}' + '__' + image_name

                shutil.copy(os.path.join(real_video_crop, image_name), os.path.join(args.dst_dir, dst_real_image))
                shutil.copy(os.path.join(fake_video_crop, image_name), os.path.join(args.dst_dir, dst_fake_image))
                f.write('{},{}\n'.format(dst_real_image, dst_fake_image))


def get_test_videos(args) -> list:
    test_videos = []
    lines = open(os.path.join(args.src_dir, LIST_OF_TESTING_VIDEOS), 'r').readlines()
    for line in lines:
        line_split = line.split()
        assert len(line_split) == 2
        if int(line_split[0]) == 1:
            continue
        video_path = os.path.join(args.src_dir, line_split[1])
        if not os.path.exists(video_path):
            continue
        video_name = os.path.basename(video_path)
        test_videos.append(video_name)
    return test_videos


def main():
    args = parse_args()
    args.src_dir = os.path.join(args.src_dir, 'Celeb-DF-v2')
    args.dst_dir = os.path.join(args.dst_dir, 'celeb-df')
    os.makedirs(args.dst_dir, exist_ok=True)

    videos = get_test_videos(args)
    np.random.seed(111)
    sampled_videos = np.random.choice(videos, 100)

    copy_five_frame(sampled_videos, args)


if __name__ == '__main__':
    main()
